package com.sui.bigdata.sml.model.configuration;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.TypeReference;
import com.feidee.fd.sml.algorithm.forecast.Forecast;
import com.feidee.fdcommon.configuration.CustomConfiguration;
import com.sui.bigdata.sml.model.exception.BorrowModelException;
import com.sui.bigdata.sml.model.model.Bootstrap;
import com.sui.bigdata.sml.model.model.Model;
import com.sui.bigdata.sml.model.util.Constants;
import com.sui.bigdata.sml.model.util.ModelFactory;
import com.sui.creep.util.StringUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.pool2.KeyedPooledObjectFactory;
import org.apache.commons.pool2.impl.GenericKeyedObjectPool;
import org.apache.commons.pool2.impl.GenericKeyedObjectPoolConfig;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.DependsOn;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @author songhaicheng
 * @date 2019/12/4 15:39
 * @description
 * @reviewer
 */
@Slf4j
@DependsOn("disconfConfig")
@Configuration
public class ModelConfig {

    private List<Model> models;

    public List<Model> getModels() {
        return models;
    }

    public ModelConfig() throws Exception {
        Bootstrap bootstrap = JSON.parseObject(CustomConfiguration.getString(Constants.ARG_BOOTSTRAP, "{}"), new TypeReference<Bootstrap>() {});
        String servingToken = bootstrap.getServingToken();
        CustomConfiguration.setString(Constants.SERVING_TOKEN, servingToken);
        models = bootstrap.getModels();
        if (StringUtils.isEmpty(servingToken) || CollectionUtils.isEmpty(models)) {
            System.out.println("模型信息获取异常");
            System.exit(0);
        }
        GenericKeyedObjectPoolConfig<Forecast> config = new GenericKeyedObjectPoolConfig();
        config.setMaxTotal(CustomConfiguration.getInt(getModelObjectPoolKey(servingToken, Constants.MODEL_POOL_MAX_TOTAL), 500));
        config.setMaxIdlePerKey(CustomConfiguration.getInt(getModelObjectPoolKey(servingToken, Constants.MODEL_POOL_MAX_IDLE_PER_KEY),10));
        config.setMaxTotalPerKey(CustomConfiguration.getInt(getModelObjectPoolKey(servingToken, Constants.MODEL_POOL_MAX_TOTAL_PER_KEY), 50));
        config.setMinIdlePerKey(CustomConfiguration.getInt(getModelObjectPoolKey(servingToken, Constants.MODEL_POOL_MIN_IDLE_PER_KEY), 1));
        config.setBlockWhenExhausted(CustomConfiguration.getBoolean(getModelObjectPoolKey(servingToken, Constants.MODEL_POOL_BLOCK_WHEN_EXHAUSTED), true));
        config.setFairness(CustomConfiguration.getBoolean(getModelObjectPoolKey(servingToken, Constants.MODEL_POOL_FAIRNESS), true));
        config.setMaxWaitMillis(CustomConfiguration.getLong(getModelObjectPoolKey(servingToken, Constants.MODEL_POOL_BORROW_MAX_WAIT_MILLIS), 200));
        KeyedPooledObjectFactory<Model, Forecast> factory = new ModelFactory();
        log.info("本次模型servingToken：" + servingToken + "，预测对象池配置如下："
                + "\n" + "maxTotal：" + CustomConfiguration.getInt(getModelObjectPoolKey(servingToken, Constants.MODEL_POOL_MAX_TOTAL), 500)
                + "\n" + "maxIdlePerKey：" + CustomConfiguration.getInt(getModelObjectPoolKey(servingToken, Constants.MODEL_POOL_MAX_IDLE_PER_KEY), 10)
                + "\n" + "maxTotalPerKey：" + CustomConfiguration.getInt(getModelObjectPoolKey(servingToken, Constants.MODEL_POOL_MAX_TOTAL_PER_KEY), 50)
                + "\n" + "minIdlePerKey：" + CustomConfiguration.getInt(getModelObjectPoolKey(servingToken, Constants.MODEL_POOL_MIN_IDLE_PER_KEY), 1)
                + "\n" + "blockWhenExhausted：" + CustomConfiguration.getBoolean(getModelObjectPoolKey(servingToken, Constants.MODEL_POOL_BLOCK_WHEN_EXHAUSTED), true)
                + "\n" + "fairness：" + CustomConfiguration.getBoolean(getModelObjectPoolKey(servingToken, Constants.MODEL_POOL_FAIRNESS), true)
                + "\n" + "borrowMaxWaitMillis：" + CustomConfiguration.getLong(getModelObjectPoolKey(servingToken, Constants.MODEL_POOL_BORROW_MAX_WAIT_MILLIS), 200)
        );
        pool = new GenericKeyedObjectPool(factory, config);
        justModel = new HashMap<>();
        for (Model model : models) {
            if (model.isSingleton()) {
                justModel.put(model, new ModelFactory().create(model));
            } else {
                pool.preparePool(model);
            }
        }

        CustomConfiguration.setString(
                Constants.LOG_MODELS,
                models.parallelStream().map(Model::getToken).collect(Collectors.joining(","))
        );
        CustomConfiguration.setString(
                Constants.LOG_MODEL_PATHS,
                models.parallelStream().map(Model::getHdfsPath).collect(Collectors.joining(", "))
        );
        log.info("=========== 模型加载完成：{} ===========", CustomConfiguration.getString(Constants.LOG_MODEL_PATHS));
    }

    /**
     * 模型池
     */
    private static GenericKeyedObjectPool<Model, Forecast> pool;
    /**
     * 单例表
     */
    private static Map<Model, Forecast> justModel;

    /**
     * 获取预测对象
     *
     * @param model
     * @return
     */
    public Forecast borrowModel(Model model) {
        Forecast forecast;
        if (model.isSingleton()) {
            forecast = justModel.get(model);
        } else {
            long start = System.currentTimeMillis();
            try {
                forecast = pool.borrowObject(model);
            } catch (Exception e) {
                throw new BorrowModelException(model);
            }
            log.debug("borrow model cost {}ms", (System.currentTimeMillis() - start));
        }
        return forecast;
    }

    /**
     * 归还预测对象
     *
     * @param model
     * @param forecast
     */
    public void returnModel(Model model, Forecast forecast) {
        if (forecast != null && !model.isSingleton()) {
            long start = System.currentTimeMillis();
            pool.returnObject(model, forecast);
            log.debug("return model cost {}ms", (System.currentTimeMillis() - start));
        }
    }

    /*
     * @Author: chq
     * @desc: 封装模型池配置项Name
     * @date: 2020/11/30 10:29
     */
    public String getModelObjectPoolKey(String servingToken, String poolConfigName) {
        return servingToken + "."+poolConfigName;
    }
}
